import pandas as pd
import matplotlib.pyplot as plt
import plotly.express as px
import numpy as np

# 设置中文字体支持
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


# 实验1.1 搭建环境并读取数据
def load_data():
    """加载并返回人口、面积和缩写三个数据集"""
    try:
        # 从URL加载示例数据
        pop_url = 'https://raw.githubusercontent.com/jakevdp/data-USstates/master/state-population.csv'
        areas_url = 'https://raw.githubusercontent.com/jakevdp/data-USstates/master/state-areas.csv'
        abbrevs_url = 'https://raw.githubusercontent.com/jakevdp/data-USstates/master/state-abbrevs.csv'

        pop = pd.read_csv(pop_url)
        areas = pd.read_csv(areas_url)
        abbrevs = pd.read_csv(abbrevs_url)

        print("数据加载成功！")
        print("\n人口数据前5行：")
        print(pop.head())
        print("\n面积数据前5行：")
        print(areas.head())
        print("\n缩写数据前5行：")
        print(abbrevs.head())

        return pop, areas, abbrevs

    except Exception as e:
        print(f"数据加载失败: {e}")
        return None, None, None


# 实验1.2 列合并
def merge_data(pop, areas, abbrevs):
    """合并三个数据集并处理缺失值"""
    if pop is None or areas is None or abbrevs is None:
        return None

    # 合并人口数据和缩写数据
    merged = pd.merge(pop, abbrevs, how='outer',
                      left_on='state/region', right_on='abbreviation')
    merged = merged.drop('abbreviation', 1)  # 丢弃重复列

    print("\n合并后的数据前3行：")
    print(merged.head(3))

    # 处理缺失值
    print("\n检查缺失值：")
    print(merged.isnull().any())

    # 查看哪些州有缺失值
    print("\n州名缺失的地区：")
    print(merged.loc[merged['state'].isnull(), 'state/region'].unique())

    # 填充缺失的州名
    merged.loc[merged['state/region'] == 'PR', 'state'] = 'Puerto Rico'
    merged.loc[merged['state/region'] == 'USA', 'state'] = 'United States'

    # 合并面积数据
    final = pd.merge(merged, areas, on='state', how='left')

    # 检查合并后的缺失值
    print("\n合并面积数据后检查缺失值：")
    print(final.isnull().any())

    # 查看面积缺失的地区
    print("\n面积缺失的州：")
    print(final['state'][final['area (sq. mi)'].isnull()].unique())

    # 删除包含缺失值的行
    final.dropna(inplace=True)

    return final


# 实验1.3 查看缺失数据、设置索引、排序
def analyze_data(final):
    """分析数据并计算人口密度"""
    if final is None:
        return None

    # 筛选2010年总人口数据
    data2010 = final.query("year == 2010 & ages == 'total'")

    # 设置州名为索引
    data2010.set_index('state', inplace=True)

    # 计算人口密度（人/平方英里）
    data2010['density'] = data2010['population'] / data2010['area (sq. mi)']

    # 按人口密度排序
    density_sorted = data2010['density'].sort_values(ascending=False)

    print("\n人口密度最高的5个州：")
    print(density_sorted.head())

    print("\n人口密度最低的5个州：")
    print(density_sorted.tail())

    return data2010, density_sorted


# 实验1.4 数据可视化
def visualize_data(data2010, density_sorted):
    """使用matplotlib和plotly进行数据可视化"""
    if data2010 is None or density_sorted is None:
        return

    # 1. 使用matplotlib绘制人口密度最高的10个州
    plt.figure(figsize=(12, 6))
    density_sorted.head(10).plot(kind='bar')
    plt.title('2010年美国人口密度最高的10个州')
    plt.xlabel('州名')
    plt.ylabel('人口密度 (人/平方英里)')
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.savefig('top10_density.png')
    plt.show()

    # 2. 使用plotly绘制美国地图上的人口密度分布
    # 准备地图数据
    map_data = data2010.reset_index()
    # 过滤掉非州的数据（如华盛顿特区）
    states_data = map_data[~map_data['state'].isin(['District of Columbia', 'Puerto Rico'])]

    fig = px.choropleth(
        states_data,
        locations='state',
        locationmode='USA-states',
        color='density',
        scope='usa',
        color_continuous_scale='Viridis',
        title='2010年美国各州人口密度分布',
        labels={'density': '人口密度 (人/平方英里)'}
    )

    fig.show()

    # 3. 绘制人口与面积的散点图
    plt.figure(figsize=(10, 6))
    plt.scatter(data2010['area (sq. mi)'], data2010['population'],
                alpha=0.7, s=data2010['density'] / 5, c=data2010['density'],
                cmap='viridis')
    plt.colorbar(label='人口密度')
    plt.title('2010年美国各州人口与面积关系')
    plt.xlabel('面积 (平方英里)')
    plt.ylabel('人口')

    # 添加几个州的标签作为示例
    for state in ['California', 'Texas', 'New York', 'Alaska', 'New Jersey']:
        plt.annotate(state,
                     (data2010.loc[state, 'area (sq. mi)'],
                      data2010.loc[state, 'population']),
                     textcoords="offset points",
                     xytext=(0, 10),
                     ha='center')

    plt.tight_layout()
    plt.savefig('population_area_scatter.png')
    plt.show()


if __name__ == "__main__":
    # 执行实验1的完整流程
    print("开始执行实验1：美国各州的统计数据\n")

    # 步骤1：加载数据
    print("\n=== 步骤1.1 加载数据 ===")
    pop, areas, abbrevs = load_data()

    # 步骤2：合并数据
    print("\n=== 步骤1.2 合并数据 ===")
    final = merge_data(pop, areas, abbrevs)

    # 步骤3：分析数据
    print("\n=== 步骤1.3 分析数据 ===")
    data2010, density_sorted = analyze_data(final)
    # 步骤4：可视化数据
    print("\n=== 步骤1.4 可视化数据 ===")
    visualize_data(data2010, density_sorted)
    print("\n实验1执行完毕！")
